Introduce `expr_fields` to `AccumulatorArgs` to hold input argument fields #18100

Jefffrey · 2025-10-16T15:23:46Z

Which issue does this PR close?

Rationale for this change

When reviewing #17085 I was very confused by the fix suggested, and tried to understand why AccumulatorArgs didn't have easy access to Fields of its input expressions, as compared to scalar/window functions which do. Introducing this new field should make it easier for users to grab datatype, metadata, nullability of their input expressions for aggregate functions.

What changes are included in this PR?

Add a slice of FieldRef to AccumulatorArgs so users don't need to compute the input expression fields themselves via using schema. This addresses #16997 as it was confusing to have only the schema available as there are valid (?) cases where the schema is empty (such as literal only input).

This fix differs from #17085 in that it doesn't special case for when there is literal only input; it leaves the physical schema provided to AccumulatorArgs untouched but provides a more ergonomic (and less confusing) API for users to retrieve Fields of their input arguments.

I'm still not sure if the schema being empty for literal only inputs is correct or not, so this might be considered a side step. If we could remove schema entirely from AccumulatorArgs maybe we wouldn't need to worry about this, but see my comment for why that wasn't done in this PR

Are these changes tested?

Existing unit tests.

Are there any user-facing changes?

Yes, new field to AccumulatorArgs which is publicly exposed (with all it's fields).

…ields

Jefffrey

I wanted to remove schema entirely, to follow what #11725 aims for. However there are some usages that I couldn't trivially fix:

datafusion/datafusion/functions-aggregate/src/array_agg.rs

Lines 207 to 210 in 9bfa2ae

    
           let ordering_dtypes = ordering 
        
               .iter() 
        
               .map(|e| e.expr.data_type(acc_args.schema)) 
        
               .collect::<Result<Vec<_>>>()?;

datafusion/datafusion/functions-aggregate/src/nth_value.rs

Lines 158 to 161 in 9bfa2ae

    
           let ordering_dtypes = ordering 
        
               .iter() 
        
               .map(|e| e.expr.data_type(acc_args.schema)) 
        
               .collect::<Result<Vec<_>>>()?;

datafusion/datafusion/functions-aggregate/src/first_last.rs

Lines 148 to 151 in 9bfa2ae

    
           let ordering_dtypes = ordering 
        
               .iter() 
        
               .map(|e| e.expr.data_type(acc_args.schema)) 
        
               .collect::<Result<Vec<_>>>()?;

Not to mention it might be more breaking to remove it (we could deprecate it I guess).

Jefffrey · 2025-10-16T15:25:46Z

datafusion/functions-aggregate-common/src/accumulator.rs

+    /// Fields corresponding to each expr (same order & length).
+    pub expr_fields: &'a [FieldRef],


Main change is here

Jefffrey · 2025-10-16T15:27:06Z

datafusion/physical-expr/src/aggregate.rs

+        let arg_fields = args
+            .iter()
+            .map(|e| e.return_field(schema.as_ref()))
+            .collect::<Result<Vec<_>>>()?;


Here is how we construct the eventual expr_fields; we're essentially doing this e.return_field(schema) pattern up front for the user, instead of requiring them to do it each time they need the field (see the various other fixes in this PR which is replacing those accesses with a simplified version)

Jefffrey · 2025-10-16T15:47:36Z

fyi @kosiew

I tried implementing like this and it seems like no issues with regressions, thoughts on if this fix is simpler?

kosiew · 2025-10-17T01:38:16Z

@Jefffrey

Your approach is an improvement!

✅ Simpler implementation - straightforward addition of pre-computed fields
✅ Less cognitive overhead - users don't need to understand schema synthesis
✅ Consistent API - same approach works for all cases (columns, literals, mixed)
✅ Performance - fields computed once during construction, not per-accumulator
✅ Explicit contract - makes it clear that field metadata is always available

👍 👍 👍

kosiew

hi @Jefffrey ,

I think you missed these:

https://github.com/apache/datafusion/blob/9bfa2ae/datafusion/physical-expr/src/aggregate.rs#L616

When AggregateFunctionExpr::with_new_expressions rewrites an aggregate to use new argument expressions, it clones the prior arg_fields. This means the expr_fields handed to the accumulator can carry metadata from the old expressions instead of the newly supplied ones, defeating the purpose of the added field (e.g., extension metadata will never update for rewritten expressions).

https://github.com/apache/datafusion/blob/9bfa2ae/datafusion/functions-aggregate/src/string_agg.rs#L194-L203

https://github.com/apache/datafusion/blob/9bfa2ae/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs#L204-L224

Both approx_percentile_cont_with_weight and the ordered/distinct branch of string_agg build nested AccumulatorArgs by filtering acc_args.exprs, but they reuse the original expr_fields slice unchanged. After the filtering, the positions (and lengths) of exprs and expr_fields no longer match, so downstream code will read the wrong field metadata once it looks beyond the first argument.

Jefffrey · 2025-10-17T05:59:21Z

hi @Jefffrey ,

I think you missed these: https://github.com/Jefffrey/datafusion/blob/acc_args_input_fields/datafusion/physical-expr/src/aggregate.rs#L625-L666

When AggregateFunctionExpr::with_new_expressions rewrites an aggregate to use new argument expressions, it clones the prior arg_fields. This means the expr_fields handed to the accumulator can carry metadata from the old expressions instead of the newly supplied ones, defeating the purpose of the added field (e.g., extension metadata will never update for rewritten expressions).

https://github.com/Jefffrey/datafusion/blob/acc_args_input_fields/datafusion/functions-aggregate/src/string_agg.rs#L194-L203

https://github.com/Jefffrey/datafusion/blob/acc_args_input_fields/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs#L204-L224

Both approx_percentile_cont_with_weight and the ordered/distinct branch of string_agg build nested AccumulatorArgs by filtering acc_args.exprs, but they reuse the original expr_fields slice unchanged. After the filtering, the positions (and lengths) of exprs and expr_fields no longer match, so downstream code will read the wrong field metadata once it looks beyond the first argument.

Thanks for picking up on this; I actually did try doing the change inside AggregateFunctionExpr::with_new_expressions but ran into some issues with async UDFs. I forgot to consider the metadata would change too. I'll look into this and see if I can understand what's going on.

Jefffrey · 2025-10-18T01:52:06Z

1c87652 fixes the AccumularArgs reconstruction for approx_percentile_cont_with_weight and string_agg.

Still looking into AggregateFunctionExpr::with_new_expressions

Jefffrey · 2025-10-18T06:34:40Z

I think there is a wider issue with with_new_expressions(); I've raised #18149 for it and left a comment in this PR, hopefully that is sufficient @kosiew ?

kosiew

LGTM

alamb · 2025-10-23T20:45:50Z

datafusion/ffi/src/udaf/accumulator_args.rs

 pub struct ForeignAccumulatorArgs {
    pub return_field: FieldRef,
    pub schema: Schema,
+    pub expr_fields: Vec<FieldRef>,


will this be a breaking FFI change?

I don't think so, since FFI_AccumulatorArgs seems to be the one marked as being stable across FFI boundaries:

datafusion/datafusion/ffi/src/udaf/accumulator_args.rs

Lines 44 to 56 in 9f23680

/// A stable struct for sharing [`AccumulatorArgs`] across FFI boundaries.

/// For an explanation of each field, see the corresponding field

/// defined in [`AccumulatorArgs`].

#[repr(C)]

#[derive(Debug, StableAbi)]

#[allow(non_camel_case_types)]

pub struct FFI_AccumulatorArgs {

return_field: WrappedSchema,

schema: WrappedSchema,

is_reversed: bool,

name: RString,

physical_expr_def: RVec<u8>,

}

Though I am not familiar with the FFI related code.

FYI @timsaucer -- perhaps you can confirm this doesn't mess up the FFI API

Introduce expr_fields to AccumulatorArgs to hold input argument f…

189e231

…ields

github-actions bot added physical-expr Changes to the physical-expr crates core Core DataFusion crate functions Changes to functions implementation ffi Changes to the ffi crate labels Oct 16, 2025

Jefffrey commented Oct 16, 2025

View reviewed changes

Jefffrey marked this pull request as ready for review October 16, 2025 15:47

Jefffrey added the api change Changes the API exposed to users of the crate label Oct 16, 2025

kosiew requested changes Oct 17, 2025

View reviewed changes

kosiew mentioned this pull request Oct 17, 2025

Support Schema Field Metadata in User-Defined Aggregate Functions (UDAFs) #17085

Closed

Merge branch 'main' into acc_args_input_fields

14016fe

Jefffrey marked this pull request as draft October 17, 2025 13:25

Jefffrey added 3 commits October 17, 2025 23:38

Merge branch 'main' into acc_args_input_fields

c2fd4fc

Fix reconstruction of accargs for string_agg and approx percentile

1c87652

Merge branch 'main' into acc_args_input_fields

221e67f

Jefffrey marked this pull request as ready for review October 18, 2025 06:32

comment

4c5a7b0

kosiew approved these changes Oct 20, 2025

View reviewed changes

alamb reviewed Oct 23, 2025

View reviewed changes

	let ordering_dtypes = ordering
	.iter()
	.map(\|e\| e.expr.data_type(acc_args.schema))
	.collect::<Result<Vec<_>>>()?;

		/// Fields corresponding to each expr (same order & length).
		pub expr_fields: &'a [FieldRef],

	/// A stable struct for sharing [`AccumulatorArgs`] across FFI boundaries.
	/// For an explanation of each field, see the corresponding field
	/// defined in [`AccumulatorArgs`].
	#[repr(C)]
	#[derive(Debug, StableAbi)]
	#[allow(non_camel_case_types)]
	pub struct FFI_AccumulatorArgs {
	return_field: WrappedSchema,
	schema: WrappedSchema,
	is_reversed: bool,
	name: RString,
	physical_expr_def: RVec<u8>,
	}

Introduce expr_fields to AccumulatorArgs to hold input argument fields #18100

Are you sure you want to change the base?

Introduce expr_fields to AccumulatorArgs to hold input argument fields #18100

Conversation

Jefffrey commented Oct 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Jefffrey left a comment

Choose a reason for hiding this comment

Uh oh!

Jefffrey Oct 16, 2025

Choose a reason for hiding this comment

Uh oh!

Jefffrey Oct 16, 2025

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Oct 16, 2025

Uh oh!

kosiew commented Oct 17, 2025

Uh oh!

kosiew left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jefffrey commented Oct 17, 2025

Uh oh!

Jefffrey commented Oct 18, 2025

Uh oh!

Jefffrey commented Oct 18, 2025

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

alamb Oct 23, 2025

Choose a reason for hiding this comment

Uh oh!

Jefffrey Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

alamb Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Introduce `expr_fields` to `AccumulatorArgs` to hold input argument fields #18100

Introduce `expr_fields` to `AccumulatorArgs` to hold input argument fields #18100

Jefffrey commented Oct 16, 2025 •

edited

Loading

kosiew left a comment •

edited

Loading

alamb Oct 24, 2025 •

edited

Loading